import json
import os

import numpy

os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"

import operator

from torch.utils.data import Dataset
import pandas as pd
import torch
import numpy as np
from vedo import *
from scipy.spatial import distance_matrix

class_2_label_lower = [
    [240, 0, 0], [251, 255, 3], [44, 251, 255], [241, 47, 255], [125, 255, 155], [26, 125, 255], [255, 234, 157],
    [206, 129, 212], [45, 135, 66], [185, 207, 45], [69, 147, 207], [207, 72, 104], [4, 207, 4], [35, 1, 207]
]

class_2_label_upper = [
    [170, 255, 127], [170, 255, 255], [255, 255, 0], [255, 170, 0], [170, 170, 255], [0, 170, 255], [85, 170, 0],
    [255, 85, 255], [255, 85, 127], [85, 170, 127], [255, 85, 0], [0, 85, 255], [170, 0, 0], [73, 247, 235]
]

def color_2_label(colors, mod):
    length = colors.shape[0]
    labels = numpy.zeros([length, 1])
    if mod == 'lower':
        cast = class_2_label_lower
    elif mod == 'upper':
        cast = class_2_label_upper
    for i in range(length):
        color = colors[i]
        for j in range(14):
            if operator.eq(color[0:3].tolist(), cast[j]):
                labels[i] = j + 1
                break
    return labels.astype('int32')


class Mesh_Dataset(Dataset):
    def __init__(self, data_list_path, num_classes=15, patch_size=6000):
        self.data_list = pd.read_csv(data_list_path, header=None)
        self.num_classes = num_classes
        self.patch_size = patch_size

    def __len__(self):
        return self.data_list.shape[0]

    def __getitem__(self, idx):
        with open('config.json', 'r') as f:
            config = json.load(f)

        mod = config["mod"]
        label_path = config["data"]["label_data_path"] + mod

        if torch.is_tensor(idx):
            idx = idx.tolist()

        i_mesh = self.data_list.iloc[idx][0]
        label_i_mesh = os.path.join(label_path, i_mesh.split('_')[-1])

        mesh = Mesh(i_mesh)
        label_mesh = Mesh(label_i_mesh)

        colors = mesh.cellcolors[:, 0:3] / 255.
        label_colors = label_mesh.cellcolors
        labels = color_2_label(label_colors, mod=mod)

        points = mesh.points()
        mean_cell_centers = mesh.center_of_mass()
        points[:, 0:3] -= mean_cell_centers[0:3]

        ids = np.array(mesh.faces())
        cells = points[ids].reshape(mesh.ncells, 9).astype(dtype='float32')

        mesh.compute_normals()
        normals = mesh.celldata['Normals']

        barycenters = mesh.cell_centers()
        barycenters -= mean_cell_centers[0:3]

        maxs = points.max(axis=0)
        mins = points.min(axis=0)
        means = points.mean(axis=0)
        stds = points.std(axis=0)
        nmeans = normals.mean(axis=0)
        nstds = normals.std(axis=0)

        for i in range(3):
            cells[:, i] = (cells[:, i] - means[i]) / stds[i]
            cells[:, i + 3] = (cells[:, i + 3] - means[i]) / stds[i]
            cells[:, i + 6] = (cells[:, i + 6] - means[i]) / stds[i]
            barycenters[:, i] = (barycenters[:, i] - mins[i]) / (maxs[i] - mins[i])
            normals[:, i] = (normals[:, i] - nmeans[i]) / nstds[i]

        X = np.column_stack((cells, barycenters, normals, colors))
        Y = labels

        X_train = np.zeros([self.patch_size, X.shape[1]], dtype='float32')
        Y_train = np.zeros([self.patch_size, Y.shape[1]], dtype='int32')
        S1 = np.zeros([self.patch_size, self.patch_size], dtype='float32')
        S2 = np.zeros([self.patch_size, self.patch_size], dtype='float32')

        valid_idx = np.argwhere(labels>=0)[:, 0]
        assert len(valid_idx) > self.patch_size

        selected_idx = np.random.choice(valid_idx, size=self.patch_size, replace=False)
        selected_idx = np.sort(selected_idx, axis=None)

        X_train[:] = X[selected_idx, :]
        Y_train[:] = Y[selected_idx, :]

        if torch.cuda.is_available():
            TX = torch.as_tensor(X_train[:, 9:12], device='cuda')
            TD = torch.cdist(TX, TX)
            D = TD.cpu().numpy()
        else:
            D = distance_matrix(X_train[:, 9:12], X_train[:, 9:12])

        S1[D<0.1] = 1.0
        S1 = S1 / np.dot(np.sum(S1, axis=1, keepdims=True), np.ones((1, self.patch_size)))

        S2[D<0.2] = 1.0
        S2 = S2 / np.dot(np.sum(S2, axis=1, keepdims=True), np.ones((1, self.patch_size)))

        X_train = X_train.transpose(1, 0)
        Y_train = Y_train.transpose(1, 0)

        sample = {'cells': torch.from_numpy(X_train),
                  'labels': torch.from_numpy(Y_train),
                  'A_S': torch.from_numpy(S1),
                  'A_L': torch.from_numpy(S2)}

        return sample


if __name__ == '__main__':
    dataset = Mesh_Dataset('./list/train_list_1.csv')
    sample = dataset.__getitem__(0)
    cells = sample["cells"]
    labels = sample["labels"]
    # for cell in cells:
    #     print(cell)
    labelnum = [0 for i in range(0, 16)]
    label = labels[0]
    print(label)
    for cal in label:
        labelnum[cal] += 1
    print(labelnum)
